{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f27c0e174c8ee037",
   "metadata": {},
   "source": [
    "# Burr\n",
    "Burr is the best library for implementing the flow chart mental model for agents.\n",
    "\n",
    "# How it works\n",
    "\n",
    "### 1. More like FastAPI in approach than LlamaIndex, LangChain, etc.\n",
    "Good at what it does so you don’t have to be.\n",
    "\n",
    "### 2. You can use whatever you like inside\n",
    "This is where you need to be good at what you do; e.g. use haystack, hamilton, langchain, even Burr itself!\n",
    "\n",
    "### 3. No need to instrument\n",
    "The framework does a lot of the work for you. E.g. opentelemetry, state checkpoints, etc.\n",
    "\n",
    "### 4. Annotation & curation\n",
    "The UI allows you to annotate what has been observed.\n",
    "\n",
    "### 5. Creating test cases, restarting & forking state\n",
    "Burr's integrated observability allows you to easily create test cases from prior executions, as well as more easily replay and debug state + code.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67546f2c-4c95-47e2-8184-234d96dc4121",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install burr[start,opentelemetry] opentelemetry-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b5c3a56971245bf",
   "metadata": {},
   "source": [
    "# Creating Actions \n",
    "A simple counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6ef1ce9e7abfbbd7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T18:22:43.114563Z",
     "start_time": "2024-10-29T18:22:42.001295Z"
    }
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "\n",
    "from burr.core import action, State\n",
    "\n",
    "\n",
    "@action(reads=[\"counter\"], writes=[\"counter\"])\n",
    "def increment(state: State) -> State:\n",
    "    \"\"\"Increment the counter by 1\"\"\"\n",
    "    current_count = state[\"counter\"]  # get the value from the `state`\n",
    "    current_count += 1\n",
    "    print(\"Count: \", current_count)\n",
    "    # use `.update()` to create a new `State`\n",
    "    return state.update(counter=current_count)\n",
    "\n",
    "\n",
    "@action(reads=[\"counter\"], writes=[])\n",
    "def exit_counter(state: State) -> State:\n",
    "    \"\"\"Print the current count and the current time\"\"\"\n",
    "    current_count = state[\"counter\"]\n",
    "    print(f\"Finished counting to {current_count} at {datetime.datetime.now():%H:%M:%S %Y-%m-%d}\")\n",
    "    return state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f18b6d1612507cb4",
   "metadata": {},
   "source": [
    "# Creating the Graph / flow chart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2d0cda40261911",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T18:25:47.592570Z",
     "start_time": "2024-10-29T18:25:46.916642Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"184pt\" height=\"112pt\"\n",
       " viewBox=\"0.00 0.00 184.10 112.20\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 108.2)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-108.2 180.1,-108.2 180.1,4 -4,4\"/>\n",
       "<!-- increment -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>increment</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M76.35,-104.2C76.35,-104.2 18.75,-104.2 18.75,-104.2 12.75,-104.2 6.75,-98.2 6.75,-92.2 6.75,-92.2 6.75,-79.6 6.75,-79.6 6.75,-73.6 12.75,-67.6 18.75,-67.6 18.75,-67.6 76.35,-67.6 76.35,-67.6 82.35,-67.6 88.35,-73.6 88.35,-79.6 88.35,-79.6 88.35,-92.2 88.35,-92.2 88.35,-98.2 82.35,-104.2 76.35,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.55\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">increment</text>\n",
       "</g>\n",
       "<!-- increment&#45;&gt;increment -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>increment&#45;&gt;increment</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M88.76,-92.14C98.85,-91.68 106.35,-89.6 106.35,-85.9 106.35,-83.82 103.98,-82.25 100.1,-81.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"100.64,-77.74 90.26,-79.86 99.7,-84.67 100.64,-77.74\"/>\n",
       "<text text-anchor=\"middle\" x=\"141.22\" y=\"-80.85\" font-family=\"Times,serif\" font-size=\"14.00\">counter &lt; 10</text>\n",
       "</g>\n",
       "<!-- exit_counter -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>exit_counter</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M83.1,-36.6C83.1,-36.6 12,-36.6 12,-36.6 6,-36.6 0,-30.6 0,-24.6 0,-24.6 0,-12 0,-12 0,-6 6,0 12,0 12,0 83.1,0 83.1,0 89.1,0 95.1,-6 95.1,-12 95.1,-12 95.1,-24.6 95.1,-24.6 95.1,-30.6 89.1,-36.6 83.1,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.55\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">exit_counter</text>\n",
       "</g>\n",
       "<!-- increment&#45;&gt;exit_counter -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>increment&#45;&gt;exit_counter</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M47.55,-67.32C47.55,-61.44 47.55,-54.72 47.55,-48.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"51.05,-48.44 47.55,-38.44 44.05,-48.44 51.05,-48.44\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "Graph(actions=[increment: counter -> counter, exit_counter: counter -> {}], transitions=[Transition(from_=increment: counter -> counter, to=increment: counter -> counter, condition=condition: counter < 10), Transition(from_=increment: counter -> counter, to=exit_counter: counter -> {}, condition=condition: default)])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from burr.core import ApplicationBuilder, expr, default\n",
    "from burr.core.graph import GraphBuilder\n",
    "graph = (\n",
    "    GraphBuilder()\n",
    "    .with_actions(increment, exit_counter)\n",
    "    .with_transitions(\n",
    "        (\"increment\", \"increment\", expr(\"counter < 10\")),\n",
    "        (\"increment\", \"exit_counter\", default),\n",
    "    ).build()\n",
    ")\n",
    "graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2d0239f76ab3397",
   "metadata": {},
   "source": [
    "# Running the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "67b4ea52feea8040",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T18:28:56.162095Z",
     "start_time": "2024-10-29T18:28:56.156315Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count:  1\n",
      "Count:  2\n",
      "Count:  3\n",
      "Count:  4\n",
      "Count:  5\n",
      "Count:  6\n",
      "Count:  7\n",
      "Count:  8\n",
      "Count:  9\n",
      "Count:  10\n",
      "Finished counting to 10 at 22:26:54 2024-10-29\n"
     ]
    }
   ],
   "source": [
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(graph)\n",
    "    .with_state(counter=0)\n",
    "    .with_entrypoint(\"increment\")\n",
    "    .build()\n",
    ")\n",
    "action_obj, result, state = app.run(halt_after=[\"exit_counter\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5aa491b5065b34b2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T18:29:08.417393Z",
     "start_time": "2024-10-29T18:29:08.414274Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "This is trying to return without having computed a single action -- we'll end up just returning some Nones. This means that nothing was executed (E.G. that the state machine had nowhere to go). Either fix the state machine orthe halt conditions, or both... Halt conditions are: halt_before=[], halt_after=['increment'].Note that this is considered undefined behavior -- if you get here, you should fix!\n"
     ]
    }
   ],
   "source": [
    "# action_obj, result, state = app.run(halt_after=[\"increment\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48fd633-87c1-4afb-928a-fd85bfaa9fc0",
   "metadata": {},
   "source": [
    "# Example extensibility - Hooks!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "06a76697-cd10-4b06-ab79-e8b22d4654ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "counter state ======>\t 1\n",
      "counter state ======>\t 2\n",
      "counter state ======>\t 3\n",
      "counter state ======>\t 4\n",
      "counter state ======>\t 5\n",
      "counter state ======>\t 6\n",
      "counter state ======>\t 7\n",
      "counter state ======>\t 8\n",
      "counter state ======>\t 9\n",
      "counter state ======>\t 10\n",
      "Finished counting to 10 at 19:15:25 2024-10-29\n",
      "counter state ======>\t 10\n"
     ]
    }
   ],
   "source": [
    "from burr.lifecycle import PostRunStepHook\n",
    "\n",
    "class PrintStepHook(PostRunStepHook):\n",
    "    \"\"\"Prints the action and state after each step.\"\"\"\n",
    "\n",
    "    def post_run_step(self, *, state: \"State\", action: \"Action\", **future_kwargs):\n",
    "        # print(\"action=====\\n\", action)\n",
    "        print(\"counter state ======>\\t\", state[\"counter\"])\n",
    "\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(graph)\n",
    "    .with_state(counter=0)\n",
    "    .with_entrypoint(\"increment\")\n",
    "    .with_hooks(PrintStepHook())\n",
    "    .build()\n",
    ")\n",
    "action_obj, result, state = app.run(halt_after=[\"exit_counter\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50c4677748eb57ba",
   "metadata": {},
   "source": [
    "# Multi-modal example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "94bef3324e93520",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T21:11:22.793863Z",
     "start_time": "2024-10-29T21:11:21.205521Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"576pt\" height=\"331pt\"\n",
       " viewBox=\"0.00 0.00 575.58 330.50\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 326.5)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-326.5 571.58,-326.5 571.58,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M509.1,-254.9C509.1,-254.9 468.75,-254.9 468.75,-254.9 462.75,-254.9 456.75,-248.9 456.75,-242.9 456.75,-242.9 456.75,-230.3 456.75,-230.3 456.75,-224.3 462.75,-218.3 468.75,-218.3 468.75,-218.3 509.1,-218.3 509.1,-218.3 515.1,-218.3 521.1,-224.3 521.1,-230.3 521.1,-230.3 521.1,-242.9 521.1,-242.9 521.1,-248.9 515.1,-254.9 509.1,-254.9\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-230.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M433.6,-187.3C433.6,-187.3 354.25,-187.3 354.25,-187.3 348.25,-187.3 342.25,-181.3 342.25,-175.3 342.25,-175.3 342.25,-162.7 342.25,-162.7 342.25,-156.7 348.25,-150.7 354.25,-150.7 354.25,-150.7 433.6,-150.7 433.6,-150.7 439.6,-150.7 445.6,-156.7 445.6,-162.7 445.6,-162.7 445.6,-175.3 445.6,-175.3 445.6,-181.3 439.6,-187.3 433.6,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"393.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M463.47,-218.02C452.82,-210.67 440.27,-202 428.79,-194.07\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"430.92,-191.29 420.7,-188.49 426.94,-197.05 430.92,-191.29\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"539.48,-322.5 438.38,-322.5 438.38,-285.9 539.48,-285.9 539.48,-322.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-298.4\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M488.93,-285.62C488.93,-279.74 488.93,-273.02 488.93,-266.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"492.43,-266.74 488.93,-256.74 485.43,-266.74 492.43,-266.74\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M117.85,-104.2C117.85,-104.2 22,-104.2 22,-104.2 16,-104.2 10,-98.2 10,-92.2 10,-92.2 10,-79.6 10,-79.6 10,-73.6 16,-67.6 22,-67.6 22,-67.6 117.85,-67.6 117.85,-67.6 123.85,-67.6 129.85,-73.6 129.85,-79.6 129.85,-79.6 129.85,-92.2 129.85,-92.2 129.85,-98.2 123.85,-104.2 117.85,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"69.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M341.99,-167.51C269.03,-165.96 140.74,-159.76 101.68,-135.7 93.45,-130.63 86.93,-122.54 82.01,-114.39\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"85.11,-112.77 77.31,-105.6 78.93,-116.07 85.11,-112.77\"/>\n",
       "<text text-anchor=\"middle\" x=\"164.55\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M247.73,-104.2C247.73,-104.2 160.12,-104.2 160.12,-104.2 154.12,-104.2 148.12,-98.2 148.12,-92.2 148.12,-92.2 148.12,-79.6 148.12,-79.6 148.12,-73.6 154.12,-67.6 160.12,-67.6 160.12,-67.6 247.73,-67.6 247.73,-67.6 253.73,-67.6 259.73,-73.6 259.73,-79.6 259.73,-79.6 259.73,-92.2 259.73,-92.2 259.73,-98.2 253.73,-104.2 247.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"203.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M342.02,-163.6C313.66,-159.45 278.75,-151.38 250.93,-135.7 240.73,-129.96 231.37,-121.31 223.72,-112.88\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"226.55,-110.81 217.4,-105.47 221.22,-115.35 226.55,-110.81\"/>\n",
       "<text text-anchor=\"middle\" x=\"309.43\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M389.73,-104.2C389.73,-104.2 290.12,-104.2 290.12,-104.2 284.12,-104.2 278.12,-98.2 278.12,-92.2 278.12,-92.2 278.12,-79.6 278.12,-79.6 278.12,-73.6 284.12,-67.6 290.12,-67.6 290.12,-67.6 389.73,-67.6 389.73,-67.6 395.73,-67.6 401.73,-73.6 401.73,-79.6 401.73,-79.6 401.73,-92.2 401.73,-92.2 401.73,-98.2 395.73,-104.2 389.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M387.8,-150.49C383.99,-140.82 378.56,-128.85 371.93,-119.2 370.51,-117.14 368.95,-115.09 367.31,-113.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"370.06,-110.91 360.78,-105.8 364.85,-115.58 370.06,-110.91\"/>\n",
       "<text text-anchor=\"middle\" x=\"444.8\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M534.23,-104.2C534.23,-104.2 431.62,-104.2 431.62,-104.2 425.62,-104.2 419.62,-98.2 419.62,-92.2 419.62,-92.2 419.62,-79.6 419.62,-79.6 419.62,-73.6 425.62,-67.6 431.62,-67.6 431.62,-67.6 534.23,-67.6 534.23,-67.6 540.23,-67.6 546.23,-73.6 546.23,-79.6 546.23,-79.6 546.23,-92.2 546.23,-92.2 546.23,-98.2 540.23,-104.2 534.23,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"482.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M445.99,-160.38C474.35,-155.05 505.36,-146.87 513.92,-135.7 519.37,-128.6 517.26,-120.53 512.21,-113.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"515.23,-111.19 506.16,-105.73 509.85,-115.66 515.23,-111.19\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M366.48,-36.6C366.48,-36.6 313.38,-36.6 313.38,-36.6 307.38,-36.6 301.38,-30.6 301.38,-24.6 301.38,-24.6 301.38,-12 301.38,-12 301.38,-6 307.38,0 313.38,0 313.38,0 366.48,0 366.48,0 372.48,0 378.48,-6 378.48,-12 378.48,-12 378.48,-24.6 378.48,-24.6 378.48,-30.6 372.48,-36.6 366.48,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M130.13,-69.77C133.1,-69.03 136.05,-68.31 138.93,-67.6 190.38,-54.97 249.33,-40.84 290.02,-31.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"290.6,-34.61 299.52,-28.89 288.98,-27.8 290.6,-34.61\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"95.85,-187.3 0,-187.3 0,-150.7 95.85,-150.7 95.85,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M38.84,-150.37C35.2,-140.89 32.75,-129.13 36.93,-119.2 37.79,-117.15 38.85,-115.16 40.06,-113.25\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"42.69,-115.57 46.15,-105.56 37.2,-111.23 42.69,-115.57\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M48.56,-150.28C49.89,-139.79 53.35,-127.06 61.93,-119.2 85.32,-97.76 101.77,-109.4 136.97,-104.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.31,-107.65 146.48,-102.34 136,-100.77 137.31,-107.65\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M96.05,-162.46C128.21,-157.75 171.03,-149.48 206.93,-135.7 220.41,-130.52 221.71,-125.03 234.93,-119.2 245.07,-114.73 256.02,-110.62 266.88,-106.94\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"267.83,-110.31 276.25,-103.88 265.66,-103.66 267.83,-110.31\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M240.73,-67.15C256.75,-59.42 275.73,-50.26 292.77,-42.04\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"294.24,-45.22 301.73,-37.72 291.2,-38.92 294.24,-45.22\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M339.93,-67.32C339.93,-61.44 339.93,-54.72 339.93,-48.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"343.43,-48.44 339.93,-38.44 336.43,-48.44 343.43,-48.44\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M444.23,-67.15C427.22,-59.34 407.05,-50.09 389,-41.81\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"390.58,-38.69 380.03,-37.7 387.66,-45.05 390.58,-38.69\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M378.9,-22.35C434.77,-27.62 533.29,-40.57 554.92,-67.6 589.49,-110.79 545,-174.31 514.04,-209.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"511.64,-207.01 507.54,-216.78 516.84,-211.7 511.64,-207.01\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "Graph(actions=[prompt: {} -> chat_history, prompt, decide_mode: prompt -> mode, generate_image: prompt, chat_history, mode -> response, generate_code: prompt, chat_history, mode -> response, answer_question: prompt, chat_history, mode -> response, prompt_for_more: prompt, chat_history -> response, response: response, mode -> chat_history], transitions=[Transition(from_=prompt: {} -> chat_history, prompt, to=decide_mode: prompt -> mode, condition=condition: default), Transition(from_=decide_mode: prompt -> mode, to=generate_image: prompt, chat_history, mode -> response, condition=condition: mode=generate_image), Transition(from_=decide_mode: prompt -> mode, to=generate_code: prompt, chat_history, mode -> response, condition=condition: mode=generate_code), Transition(from_=decide_mode: prompt -> mode, to=answer_question: prompt, chat_history, mode -> response, condition=condition: mode=answer_question), Transition(from_=decide_mode: prompt -> mode, to=prompt_for_more: prompt, chat_history -> response, condition=condition: default), Transition(from_=generate_image: prompt, chat_history, mode -> response, to=response: response, mode -> chat_history, condition=condition: default), Transition(from_=answer_question: prompt, chat_history, mode -> response, to=response: response, mode -> chat_history, condition=condition: default), Transition(from_=generate_code: prompt, chat_history, mode -> response, to=response: response, mode -> chat_history, condition=condition: default), Transition(from_=prompt_for_more: prompt, chat_history -> response, to=response: response, mode -> chat_history, condition=condition: default), Transition(from_=response: response, mode -> chat_history, to=prompt: {} -> chat_history, prompt, condition=condition: default)])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import copy\n",
    "\n",
    "import openai\n",
    "\n",
    "from burr.core import State, default, graph, when\n",
    "from burr.core.action import action\n",
    "\n",
    "MODES = {\n",
    "    \"answer_question\": \"text\",\n",
    "    \"generate_image\": \"image\",\n",
    "    \"generate_code\": \"code\",\n",
    "    \"unknown\": \"text\",\n",
    "}\n",
    "\n",
    "\n",
    "@action(reads=[], writes=[\"chat_history\", \"prompt\"])\n",
    "def process_prompt(state: State, prompt: str) -> State:\n",
    "    result = {\"chat_item\": {\"role\": \"user\", \"content\": prompt, \"type\": \"text\"}}\n",
    "    state = state.append(chat_history=result[\"chat_item\"])\n",
    "    state = state.update(prompt=prompt)\n",
    "    return state\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\"], writes=[\"mode\"])\n",
    "def choose_mode(state: State) -> State:\n",
    "    prompt = (\n",
    "        f\"You are a chatbot. You've been prompted this: {state['prompt']}. \"\n",
    "        f\"You have the capability of responding in the following modes: {', '.join(MODES)}. \"\n",
    "        \"Please respond with *only* a single word representing the mode that most accurately \"\n",
    "        \"corresponds to the prompt. Fr instance, if the prompt is 'draw a picture of a cat', \"\n",
    "        \"the mode would be 'generate_image'. If the prompt is \"\n",
    "        \"'what is the capital of France', the mode would be 'answer_question'.\"\n",
    "        \"If none of these modes apply, please respond with 'unknown'.\"\n",
    "    )\n",
    "\n",
    "    llm_result = openai.Client().chat.completions.create(\n",
    "        model=\"gpt-4\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant\"},\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "    )\n",
    "    content = llm_result.choices[0].message.content\n",
    "    mode = content.lower()\n",
    "    if mode not in MODES:\n",
    "        mode = \"unknown\"\n",
    "    result = {\"mode\": mode}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\"], writes=[\"response\"])\n",
    "def prompt_for_more(state: State) -> State:\n",
    "    result = {\n",
    "        \"response\": {\n",
    "            \"content\": \"None of the response modes I support apply to your question. \"\n",
    "                       \"Please clarify?\",\n",
    "            \"type\": \"text\",\n",
    "            \"role\": \"assistant\",\n",
    "        }\n",
    "    }\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def chat_response(\n",
    "        state: State, prepend_prompt: str, model: str = \"gpt-3.5-turbo\"\n",
    ") -> State:\n",
    "    \n",
    "    chat_history = copy.deepcopy(state[\"chat_history\"])\n",
    "    chat_history[-1][\"content\"] = f\"{prepend_prompt}: {chat_history[-1]['content']}\"\n",
    "    chat_history_api_format = [\n",
    "        {\n",
    "            \"role\": chat[\"role\"],\n",
    "            \"content\": chat[\"content\"],\n",
    "        }\n",
    "        for chat in chat_history\n",
    "    ]\n",
    "    client = openai.Client()\n",
    "    result = client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=chat_history_api_format,\n",
    "    )\n",
    "    text_response = result.choices[0].message.content\n",
    "    result = {\"response\": {\"content\": text_response, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def image_response(state: State, model: str = \"dall-e-2\") -> State:\n",
    "    \"\"\"Generates an image response to the prompt. Optional save function to save the image to a URL.\"\"\"\n",
    "    # raise ValueError(\"Demo error\")\n",
    "    client = openai.Client()\n",
    "    result = client.images.generate(\n",
    "        model=model, prompt=state[\"prompt\"], size=\"1024x1024\", quality=\"standard\", n=1\n",
    "    )\n",
    "    image_url = result.data[0].url\n",
    "    result = {\"response\": {\"content\": image_url, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"response\", \"mode\"], writes=[\"chat_history\"])\n",
    "def response(state: State) -> State:\n",
    "    # you'd do something specific here based on prior state\n",
    "    result = {\"chat_item\": state[\"response\"]}\n",
    "    return state.append(chat_history=result[\"chat_item\"])\n",
    "\n",
    "base_graph = (\n",
    "    graph.GraphBuilder()\n",
    "    .with_actions(\n",
    "        # these are the \"nodes\" \n",
    "        prompt=process_prompt,\n",
    "        decide_mode=choose_mode,\n",
    "        generate_image=image_response,\n",
    "        generate_code=chat_response.bind(\n",
    "            prepend_prompt=\"Please respond with *only* code and no other text (at all) to the following:\",\n",
    "        ),\n",
    "        answer_question=chat_response.bind(\n",
    "            prepend_prompt=\"Please answer the following question:\",\n",
    "        ),\n",
    "        prompt_for_more=prompt_for_more,\n",
    "        response=response,\n",
    "    )\n",
    "    .with_transitions(\n",
    "        # these are the edges between nodes, based on state.\n",
    "        (\"prompt\", \"decide_mode\", default),\n",
    "        (\"decide_mode\", \"generate_image\", when(mode=\"generate_image\")),\n",
    "        (\"decide_mode\", \"generate_code\", when(mode=\"generate_code\")),\n",
    "        (\"decide_mode\", \"answer_question\", when(mode=\"answer_question\")),\n",
    "        (\"decide_mode\", \"prompt_for_more\", default),\n",
    "        (\n",
    "            [\"generate_image\", \"answer_question\", \"generate_code\", \"prompt_for_more\"],\n",
    "            \"response\",\n",
    "        ),\n",
    "        (\"response\", \"prompt\", default),\n",
    "    )\n",
    "    .build()\n",
    ")\n",
    "base_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf3e317aa71daca0",
   "metadata": {},
   "source": [
    "# Adding Instrumentation & Running it\n",
    "Just a one line change..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4b10a2ab4248135f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T21:16:19.783967Z",
     "start_time": "2024-10-29T21:16:19.298927Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"576pt\" height=\"331pt\"\n",
       " viewBox=\"0.00 0.00 575.58 330.50\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 326.5)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-326.5 571.58,-326.5 571.58,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M509.1,-254.9C509.1,-254.9 468.75,-254.9 468.75,-254.9 462.75,-254.9 456.75,-248.9 456.75,-242.9 456.75,-242.9 456.75,-230.3 456.75,-230.3 456.75,-224.3 462.75,-218.3 468.75,-218.3 468.75,-218.3 509.1,-218.3 509.1,-218.3 515.1,-218.3 521.1,-224.3 521.1,-230.3 521.1,-230.3 521.1,-242.9 521.1,-242.9 521.1,-248.9 515.1,-254.9 509.1,-254.9\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-230.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M433.6,-187.3C433.6,-187.3 354.25,-187.3 354.25,-187.3 348.25,-187.3 342.25,-181.3 342.25,-175.3 342.25,-175.3 342.25,-162.7 342.25,-162.7 342.25,-156.7 348.25,-150.7 354.25,-150.7 354.25,-150.7 433.6,-150.7 433.6,-150.7 439.6,-150.7 445.6,-156.7 445.6,-162.7 445.6,-162.7 445.6,-175.3 445.6,-175.3 445.6,-181.3 439.6,-187.3 433.6,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"393.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M463.47,-218.02C452.82,-210.67 440.27,-202 428.79,-194.07\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"430.92,-191.29 420.7,-188.49 426.94,-197.05 430.92,-191.29\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"539.48,-322.5 438.38,-322.5 438.38,-285.9 539.48,-285.9 539.48,-322.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-298.4\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M488.93,-285.62C488.93,-279.74 488.93,-273.02 488.93,-266.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"492.43,-266.74 488.93,-256.74 485.43,-266.74 492.43,-266.74\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M117.85,-104.2C117.85,-104.2 22,-104.2 22,-104.2 16,-104.2 10,-98.2 10,-92.2 10,-92.2 10,-79.6 10,-79.6 10,-73.6 16,-67.6 22,-67.6 22,-67.6 117.85,-67.6 117.85,-67.6 123.85,-67.6 129.85,-73.6 129.85,-79.6 129.85,-79.6 129.85,-92.2 129.85,-92.2 129.85,-98.2 123.85,-104.2 117.85,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"69.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M341.99,-167.51C269.03,-165.96 140.74,-159.76 101.68,-135.7 93.45,-130.63 86.93,-122.54 82.01,-114.39\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"85.11,-112.77 77.31,-105.6 78.93,-116.07 85.11,-112.77\"/>\n",
       "<text text-anchor=\"middle\" x=\"164.55\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M247.73,-104.2C247.73,-104.2 160.12,-104.2 160.12,-104.2 154.12,-104.2 148.12,-98.2 148.12,-92.2 148.12,-92.2 148.12,-79.6 148.12,-79.6 148.12,-73.6 154.12,-67.6 160.12,-67.6 160.12,-67.6 247.73,-67.6 247.73,-67.6 253.73,-67.6 259.73,-73.6 259.73,-79.6 259.73,-79.6 259.73,-92.2 259.73,-92.2 259.73,-98.2 253.73,-104.2 247.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"203.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M342.02,-163.6C313.66,-159.45 278.75,-151.38 250.93,-135.7 240.73,-129.96 231.37,-121.31 223.72,-112.88\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"226.55,-110.81 217.4,-105.47 221.22,-115.35 226.55,-110.81\"/>\n",
       "<text text-anchor=\"middle\" x=\"309.43\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M389.73,-104.2C389.73,-104.2 290.12,-104.2 290.12,-104.2 284.12,-104.2 278.12,-98.2 278.12,-92.2 278.12,-92.2 278.12,-79.6 278.12,-79.6 278.12,-73.6 284.12,-67.6 290.12,-67.6 290.12,-67.6 389.73,-67.6 389.73,-67.6 395.73,-67.6 401.73,-73.6 401.73,-79.6 401.73,-79.6 401.73,-92.2 401.73,-92.2 401.73,-98.2 395.73,-104.2 389.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M387.8,-150.49C383.99,-140.82 378.56,-128.85 371.93,-119.2 370.51,-117.14 368.95,-115.09 367.31,-113.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"370.06,-110.91 360.78,-105.8 364.85,-115.58 370.06,-110.91\"/>\n",
       "<text text-anchor=\"middle\" x=\"444.8\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M534.23,-104.2C534.23,-104.2 431.62,-104.2 431.62,-104.2 425.62,-104.2 419.62,-98.2 419.62,-92.2 419.62,-92.2 419.62,-79.6 419.62,-79.6 419.62,-73.6 425.62,-67.6 431.62,-67.6 431.62,-67.6 534.23,-67.6 534.23,-67.6 540.23,-67.6 546.23,-73.6 546.23,-79.6 546.23,-79.6 546.23,-92.2 546.23,-92.2 546.23,-98.2 540.23,-104.2 534.23,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"482.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M445.99,-160.38C474.35,-155.05 505.36,-146.87 513.92,-135.7 519.37,-128.6 517.26,-120.53 512.21,-113.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"515.23,-111.19 506.16,-105.73 509.85,-115.66 515.23,-111.19\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M366.48,-36.6C366.48,-36.6 313.38,-36.6 313.38,-36.6 307.38,-36.6 301.38,-30.6 301.38,-24.6 301.38,-24.6 301.38,-12 301.38,-12 301.38,-6 307.38,0 313.38,0 313.38,0 366.48,0 366.48,0 372.48,0 378.48,-6 378.48,-12 378.48,-12 378.48,-24.6 378.48,-24.6 378.48,-30.6 372.48,-36.6 366.48,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M130.13,-69.77C133.1,-69.03 136.05,-68.31 138.93,-67.6 190.38,-54.97 249.33,-40.84 290.02,-31.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"290.6,-34.61 299.52,-28.89 288.98,-27.8 290.6,-34.61\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"95.85,-187.3 0,-187.3 0,-150.7 95.85,-150.7 95.85,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M38.84,-150.37C35.2,-140.89 32.75,-129.13 36.93,-119.2 37.79,-117.15 38.85,-115.16 40.06,-113.25\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"42.69,-115.57 46.15,-105.56 37.2,-111.23 42.69,-115.57\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M48.56,-150.28C49.89,-139.79 53.35,-127.06 61.93,-119.2 85.32,-97.76 101.77,-109.4 136.97,-104.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.31,-107.65 146.48,-102.34 136,-100.77 137.31,-107.65\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M96.05,-162.46C128.21,-157.75 171.03,-149.48 206.93,-135.7 220.41,-130.52 221.71,-125.03 234.93,-119.2 245.07,-114.73 256.02,-110.62 266.88,-106.94\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"267.83,-110.31 276.25,-103.88 265.66,-103.66 267.83,-110.31\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M240.73,-67.15C256.75,-59.42 275.73,-50.26 292.77,-42.04\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"294.24,-45.22 301.73,-37.72 291.2,-38.92 294.24,-45.22\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M339.93,-67.32C339.93,-61.44 339.93,-54.72 339.93,-48.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"343.43,-48.44 339.93,-38.44 336.43,-48.44 343.43,-48.44\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M444.23,-67.15C427.22,-59.34 407.05,-50.09 389,-41.81\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"390.58,-38.69 380.03,-37.7 387.66,-45.05 390.58,-38.69\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M378.9,-22.35C434.77,-27.62 533.29,-40.57 554.92,-67.6 589.49,-110.79 545,-174.31 514.04,-209.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"511.64,-207.01 507.54,-216.78 516.84,-211.7 511.64,-207.01\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<burr.core.application.Application at 0x120e6a9e0>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from burr.tracking import LocalTrackingClient\n",
    "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n",
    "# this will auto instrument the openAI client. No swapping of imports required!\n",
    "OpenAIInstrumentor().instrument()\n",
    "\n",
    "tracker = LocalTrackingClient(project=\"agent-demo-new\")\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker, \n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    .with_tracker(tracker, use_otel_tracing=True)  # tracking + checkpointing; one line 🪄.\n",
    "    .build()\n",
    ")\n",
    "app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8786d2399fdba8",
   "metadata": {},
   "source": [
    "# Annotation\n",
    "Let's create some data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "92a6be977ec7e58e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T21:17:13.934154Z",
     "start_time": "2024-10-29T21:16:21.614723Z"
    }
   },
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? what is the capital of france?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: The capital of France is Paris.\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? say hello world in java\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: ```java\n",
      "public class HelloWorld {\n",
      "    public static void main(String[] args) {\n",
      "        System.out.println(\"Hello, World\");\n",
      "    }\n",
      "}\n",
      "```\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "from IPython.display import Image, display \n",
    "\n",
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0191e2fa1685852",
   "metadata": {},
   "source": [
    "[Open the UI - http://localhost:7241/](http://localhost:7241)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93e33cf02a154a5",
   "metadata": {},
   "source": [
    "# Iteration, Debugging, Test cases\n",
    "\n",
    "## Replaying from a specific point in time\n",
    "\"Forking\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ad85268bf7c89f99",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T21:18:02.462542Z",
     "start_time": "2024-10-29T21:18:02.449115Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'what is the capital of france?', 'type': 'text'},\n",
       " {'content': 'The capital of France is Paris.',\n",
       "  'type': 'text',\n",
       "  'role': 'assistant'},\n",
       " {'role': 'user', 'content': 'say hello world in java', 'type': 'text'}]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show forking state.\n",
    "app_id = \"c65fdf17-6893-4b93-8aa7-d9e6fa549a18\"\n",
    "sequence_id = 4\n",
    "# partition_key = \"\"\n",
    "\n",
    "forked_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph) # this could be different...\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "        fork_from_app_id=app_id,\n",
    "        fork_from_sequence_id=sequence_id,\n",
    "        # fork_from_partition_key=partition_key\n",
    "    )\n",
    "    .with_tracker(tracker, use_otel_tracing=True)\n",
    "    .build()\n",
    ")\n",
    "forked_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fa06c46990bd6e30",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T21:18:13.195324Z",
     "start_time": "2024-10-29T21:18:04.968629Z"
    }
   },
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: ```java\n",
      "public class HelloWorld {\n",
      "    public static void main(String[] args) {\n",
      "        System.out.println(\"Hello, World!\");\n",
      "    }\n",
      "}\n",
      "```\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = forked_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44a6d4a9535c808",
   "metadata": {},
   "source": [
    "## Can create test cases using saved data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a47290467f0dc528",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-29T21:18:55.947570Z",
     "start_time": "2024-10-29T21:18:53.595654Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No state found for agent-demo-new, null, c65fdf17-6893-4b93-8aa7-d9e6fa549a180, 4.\n"
     ]
    }
   ],
   "source": [
    "!burr-test-case create  \\\n",
    "  --project-name \"agent-demo-new\" \\\n",
    "  --partition-key \"null\" \\\n",
    "  --app-id \"c65fdf17-6893-4b93-8aa7-d9e6fa549a180\" \\\n",
    "  --sequence-id 4 \\\n",
    "  --target-file-name YOUR_FIXTURE_FILE.json \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d9157d4436bdcf2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
